// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once

#include <cstdlib>
#include <string>
#include <memory>


#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/base/base64.h"

#include <boost/asio/ip/tcp.hpp>
#include <boost/beast/core.hpp>
#include <boost/beast/http.hpp>
#include <boost/beast/ssl/ssl_stream.hpp>

namespace coin2::exchange::base::feed {

class RestConnection : public std::enable_shared_from_this<RestConnection> {
 public:
  RestConnection(
    boost::asio::io_context* io_ctx,
    const std::string& host,
    int port,
    const std::string& bind_ip,
    const std::string& http_proxy)
      : io_ctx_(io_ctx),
        ssl_ctx_{boost::asio::ssl::context::sslv23},
        port_(port),
        bind_ip_(bind_ip),
        http_proxy_(http_proxy),
        resolver_(*io_ctx_) {
    // only support https
    CHECK_EQ(host.rfind(url_prefix_), 0);
    host_ = host.substr(url_prefix_.size());
    if (!http_proxy.empty()) {
      proxy_setting_ = ParseProxy(http_proxy);
    }
  }

  bool Start(boost::asio::yield_context yield) {
    return Connnect(yield);
  }

  bool Get(
      const std::string& target,
      const std::vector<std::pair<std::string, std::string>>& header,
      std::string* response,
      boost::asio::yield_context yield) {
    // try 2 times, server side close connection even keep-alive set
    boost::beast::http::status status;
    bool suppress_error = true;
    if (!Send(target, header, {}, response, &status, suppress_error, yield)) {
      Reset(yield);
      if (!Send(target, header, {}, response, &status, !suppress_error, yield)) {
        return false;
      }
    }

    return boost::beast::http::status::ok == status;
  }

  bool Post(const std::string& target, std::string* response, boost::asio::yield_context yield) {
    // try 2 times, server side close connection even keep-alive set
    boost::beast::http::status status;
    bool suppress_error = true;
    if (!Send(target, {}, {}, response, &status, suppress_error, yield, boost::beast::http::verb::post)) {
      Reset(yield);
      if (!Send(target, {}, {}, response, &status, !suppress_error, yield, boost::beast::http::verb::post)) {
        return false;
      }
    }

    return boost::beast::http::status::ok == status;
  }

  bool Send(
      const std::string& target,
      const std::vector<std::pair<std::string, std::string>>& header,
      boost::beast::http::verb method,
      const std::string& payload,
      std::string* response,
      boost::asio::yield_context yield) {
    // try 2 times, server side close connection even keep-alive set
    boost::beast::http::status status;
    bool suppress_error = true;
    if (!Send(target, header, payload, response, &status, suppress_error, yield, method)) {
      Reset(yield);
      if (!Send(target, header, payload, response, &status, !suppress_error, yield, method)) {
        return false;
      }
    }

    return boost::beast::http::status::ok == status;
  }

  bool Reset(boost::asio::yield_context yield) {
    DLOG(INFO) << "Reconnect " << host_ << " " << port_;
    return Connnect(yield);
  }

 private:

  struct ProxySetting {
    std::string host;
    std::string port;
    std::string auth;
  };

  bool Connnect(boost::asio::yield_context yield) {
    stream_ssl_.reset(new boost::beast::ssl_stream<boost::beast::tcp_stream>(*io_ctx_, ssl_ctx_));
    if (!SSL_set_tlsext_host_name(stream_ssl_->native_handle(), host_.c_str())) {
      LOG(ERROR) << "SNI error " << host_ << " " << port_;
      return false;
    }

    if (!bind_ip_.empty()) {
      boost::beast::get_lowest_layer(*stream_ssl_).socket().open(boost::asio::ip::tcp::v4());
      boost::beast::get_lowest_layer(*stream_ssl_).socket().bind(
          boost::asio::ip::tcp::endpoint(boost::asio::ip::address::from_string(bind_ip_), 0));
    }

    boost::beast::error_code ec;

    if (!http_proxy_.empty()) {
      auto const results = resolver_.async_resolve(
          proxy_setting_.host,
          proxy_setting_.port,
          yield[ec]);
      if (ec) {
        LOG(ERROR) << "resolve proxy error " << ec.message() << ", " 
            << proxy_setting_.host << ", " << proxy_setting_.port;
        return false;
      }

      bool connected = false;
      for (auto iter = results.begin(); iter != results.end(); ++iter) {
        boost::beast::get_lowest_layer(*stream_ssl_).socket().async_connect(
            *iter,
            yield[ec]);
        if (ec) {
          LOG(ERROR) << "connect proxy error " << ec.message() << ", " 
              << proxy_setting_.host << ", " << proxy_setting_.port
              << ", address " << iter->endpoint();
        } else {
          connected = true;
          break;
        }
      }
      if (!connected) {
        LOG(ERROR) << "can't connect to proxy "
            << proxy_setting_.host << ", " << proxy_setting_.port;
        return false;
      }

      std::string host = host_ + ":" + std::to_string(port_);
      boost::beast::http::request<boost::beast::http::empty_body> request(
          boost::beast::http::verb::connect,
          host,
          11);
      request.set(boost::beast::http::field::host, host);
      request.set(boost::beast::http::field::proxy_connection, "keep-alive");
      request.set(boost::beast::http::field::connection, "keep-alive");
      if (!proxy_setting_.auth.empty()) {
        request.set(boost::beast::http::field::proxy_authorization, proxy_setting_.auth);
      }
      request.prepare_payload();

      boost::beast::http::async_write(
          boost::beast::get_lowest_layer(*stream_ssl_),
          request,
          yield[ec]);
      if (ec) {
        LOG(ERROR) << "send connect to proxy error " << ec.message() << ", "
            << proxy_setting_.host << ", " << proxy_setting_.port << ","
            << request;
        return false;
      }

      boost::beast::flat_buffer buffer;
      boost::beast::http::response<boost::beast::http::empty_body> response;
      boost::beast::http::parser<false, boost::beast::http::empty_body> p;
      boost::beast::http::async_read_header(
          boost::beast::get_lowest_layer(*stream_ssl_),
          buffer,
          p,
          yield[ec]);
      if (ec) {
        LOG(ERROR) <<"wait conenct to prxoy error " << ec.message() << ", "
            << proxy_setting_.host << ", " << proxy_setting_.port << ","
            << request;
        return false;
      }
      response = std::move(p.get());
      if (response.result() != boost::beast::http::status::ok) {
        LOG(ERROR) <<"prxoy refuse " 
            << proxy_setting_.host << ", " << proxy_setting_.port << ", "
            << response;
        return false;
      } else {
         LOG(INFO) << "connect proxy success, " << response;
      }
    } else {
      // Looks up the domain name.
      auto const results = resolver_.async_resolve(
          host_,
          std::to_string(port_),
          yield[ec]);
      if (ec) {
        LOG(ERROR) << "resolve error " << ec.message() << ", " << host_ << ", " << port_;
        return false;
      }

      bool connected = false;
      for (auto iter = results.begin(); iter != results.end(); ++iter) {
        boost::beast::get_lowest_layer(*stream_ssl_).socket().async_connect(
            *iter,
            yield[ec]);
        if (ec) {
          LOG(ERROR) << "connect error " << ec.message() << ", " << host_ << ", " << port_ <<
              ", address " << iter->endpoint();
        } else {
          connected = true;
          break;
        }
      }
      if (!connected) {
        LOG(ERROR) << "can't connect to " << host_ << ", " << port_;
        return false;
      }

      boost::beast::get_lowest_layer(*stream_ssl_).expires_after(
          std::chrono::nanoseconds(timeout_ns_));
      // set option TCP_NODELAY
      boost::beast::get_lowest_layer(*stream_ssl_).socket().set_option(
          boost::asio::ip::tcp::no_delay(true));
    }

    // Performs the SSL handshake.
    stream_ssl_->async_handshake(
        boost::asio::ssl::stream_base::client,
        yield[ec]);
    if (ec) {
      LOG(ERROR) << "ssl handshake error " << ec.message() << ", " << host_ << ", " << port_;
      return false;
    }

    return true;
  }

  ProxySetting ParseProxy(const std::string& http_proxy) {
    ProxySetting proxy_setting{};
    std::string tmp = http_proxy;
    // strip protocol text if existing
    boost::replace_all(tmp, R"(https://)", "");
    boost::replace_all(tmp, R"(http://)", "");
    // retrieve proxy authrization
    if (tmp.find(R"(@)") != std::string::npos) {
      std::vector<std::string> vecs;
      boost::split(vecs, tmp, boost::is_any_of("@"));
      if (vecs.size() == 2) {
        const auto& auth = vecs[0];
        proxy_setting.auth = fmt::format(
            "Basic {}",
            ::impl::base64_encode(reinterpret_cast<const unsigned char*>(auth.c_str()), auth.size()));
        tmp = vecs[1];
      }
    }
    // retrieve addr/port
    {
      std::vector<std::string> vecs;
      boost::split(vecs, tmp, boost::is_any_of(":"));
      CHECK_EQ(vecs.size(), 2) << "Invalid http proxy " << http_proxy;
      proxy_setting.host = vecs[0];
      proxy_setting.port = vecs[1];
    }

    return proxy_setting;
  }

  bool Send(
      const std::string& target,
      const std::vector<std::pair<std::string, std::string>>& header,
      const std::string& payload,
      std::string* body,
      boost::beast::http::status* status,
      bool suppress,
      boost::asio::yield_context yield,
      boost::beast::http::verb method = boost::beast::http::verb::get) {
    boost::beast::error_code ec;

    boost::beast::http::request<boost::beast::http::string_body> request;
    const int version = 11;
    request.version(version);
    request.method(method);
    request.keep_alive(true);
    request.set(boost::beast::http::field::host, host_);
    // coinbase rest require user agent
    request.set(boost::beast::http::field::user_agent, "Beast");
    for (const auto& pr : header) {
      request.set(pr.first, pr.second);
    }
    request.target(boost::beast::string_view(target.data(), target.size()));
    if (!payload.empty()) {
      request.set("Content-Length", std::to_string(payload.size()));
      request.body() = payload;
    }
    boost::beast::http::async_write(
        *stream_ssl_,
        request,
        yield[ec]);
    if (ec) {
      if (!suppress) {
        LOG(ERROR) << "write error " << ec.message() << ", " <<
            host_ << ", " << port_ << ", " << request;
      }
      return false;
    }

    boost::beast::flat_buffer buffer;
    boost::beast::http::response<boost::beast::http::string_body> response;
    boost::beast::http::async_read(
        *stream_ssl_,
        buffer,
        response,
        yield[ec]);
    if (ec) {
      if (!suppress) {
        LOG(ERROR) <<"read error " << ec.message() << ", " <<
            host_ << ", " << port_ << ", " << request;
      }
      return false;
    }

    *status = response.result();
    *body = response.body();
    return true;
  }

  boost::asio::io_context* io_ctx_;
  boost::asio::ssl::context ssl_ctx_;
  std::string host_;
  const int port_;
  const std::string bind_ip_;
  std::unique_ptr<boost::beast::ssl_stream<boost::beast::tcp_stream>> stream_ssl_;
  std::string http_proxy_;
  ProxySetting proxy_setting_;
  boost::asio::ip::tcp::resolver resolver_;

  // Timeout config
  static constexpr int64_t timeout_ns_ = 60 * 1000000000L;
  static constexpr std::string_view url_prefix_ = "https://";
};

}  // namespace coin2::exchange::base::feed
